import ollama
import asyncio
import yaml

from config.PathConfig import OLLAMA_CONFIG

class OllamaClient:
    def __init__(self,tools = False):
        # 从 yaml 文件加载配置
        with open(OLLAMA_CONFIG, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
        if tools :
            self.model = config["ollama"]["QWEN_MODEL"]
        else:
            self.model = config["ollama"]["DS_MODEL"]
        self.host = config["ollama"]["host"]
        self.client = ollama.Client(host=self.host)
    
    async def async_chat(self, messages, tools=None):
        # 异步调用，model 从配置中读取
        return await asyncio.to_thread(
            self.client.chat,
            model=self.model,
            messages=messages,
            tools=tools
        )

# 实例化
# OllamaClient = OllamaClient()